package com.zookeeper.util;

import io.netty.util.CharsetUtil;
import lombok.SneakyThrows;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.api.ACLProvider;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.data.ACL;
import org.apache.zookeeper.data.Id;
import org.apache.zookeeper.data.Stat;
import org.apache.zookeeper.server.auth.DigestAuthenticationProvider;

import java.util.ArrayList;
import java.util.List;

/**
 * zk工具类 - 使用 Curator 客户端
 */
public class CuratorZKUtils {

    /**
     * 创建客户端（无需认证）
     * @param connectionString 连接串
     */
    public static CuratorFramework getClient(String connectionString) {
        CuratorFramework client = CuratorFrameworkFactory.builder()
                .connectString(connectionString)
                .retryPolicy(new ExponentialBackoffRetry(1000, 3))
                .build();
        client.start();
        return client;
    }

    /**
     * 创建客户端（需要认证）
     * @param scheme digest
     * @param usernameAndPassword  username:password
     * @param connectionString 连接串
     */
    @SneakyThrows
    public static CuratorFramework getClient(String connectionString, String scheme, String usernameAndPassword,
                                             String keyStoreLocal, String keyStoreLocalPW, String trustStoreLocal, String trustStoreLocalPW) {
        // 设置 SSL 连接需要的证书系统环境
        System.setProperty("zookeeper.clientCnxnSocket", "org.apache.zookeeper.ClientCnxnSocketNetty");
        System.setProperty("zookeeper.client.secure", "true");
        System.setProperty("zookeeper.ssl.keyStore.location", keyStoreLocal);
        System.setProperty("zookeeper.ssl.keyStore.password", keyStoreLocalPW);
        System.setProperty("zookeeper.ssl.trustStore.location", trustStoreLocal);
        System.setProperty("zookeeper.ssl.trustStore.password", trustStoreLocalPW);
        System.setProperty("zookeeper.ssl.hostnameVerification", "false");

        // 设置 ACL 权限
        List<ACL> aclList = new ArrayList<>();
        ACL acl = new ACL(ZooDefs.Perms.ALL, new Id(scheme,
                DigestAuthenticationProvider.generateDigest(usernameAndPassword)));
        aclList.add(acl);

        ACLProvider aclProvider = new ACLProvider() {
            @Override
            public List<ACL> getDefaultAcl() {
                return aclList;
            }

            @Override
            public List<ACL> getAclForPath(String path) {
                return aclList;
            }
        };

        CuratorFramework client = CuratorFrameworkFactory.builder()
                .connectString(connectionString)
                .authorization(scheme, usernameAndPassword.getBytes())  // 账号认证
                .aclProvider(aclProvider)
                .retryPolicy(new ExponentialBackoffRetry(1000, 3)) // 设置重试策略
                .build();
        client.start();

        // 创建客户端后就可以清空之前设置的环境
        System.clearProperty("zookeeper.clientCnxnSocket");
        System.clearProperty("zookeeper.client.secure");
        System.clearProperty("zookeeper.ssl.keyStore.location");
        System.clearProperty("zookeeper.ssl.keyStore.password");
        System.clearProperty("zookeeper.ssl.trustStore.location");
        System.clearProperty("zookeeper.ssl.trustStore.password");
        System.clearProperty("zookeeper.ssl.hostnameVerification");
        return client;
    }

    /**
     * 递归创建节点
     */
    @SneakyThrows
    public static String createNodeRecursion(CuratorFramework client, String path) {
        return client.create().creatingParentsIfNeeded().forPath(path);
    }

    /**
     * 删除节点
     */
    @SneakyThrows
    public static void delete(CuratorFramework client, String path) {
        client.delete().forPath(path);
    }

    /**
     * 获取节点数据
     */
    @SneakyThrows
    public static String getData(CuratorFramework client, String path) {
        byte[] bytes = client.getData().forPath(path);
        return new String(bytes);
    }

    /**
     * 设置节点数据
     */
    @SneakyThrows
    public static Stat setData(CuratorFramework client, String path, String value) {
        return client.setData().forPath(path, value.getBytes(CharsetUtil.UTF_8));
    }

}

